#! /usr/bin/python
# encoding=utf-8
# Created by Fenglu Niu on 2025/5/6 17:42

import json
from datetime import datetime
from typing import Annotated

from fastapi import APIRouter, Query
from fastapi.params import Depends
from langchain_core.messages import AIMessage, ToolMessage
from langchain_core.prompts import PromptTemplate, ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnablePassthrough
from sqlalchemy import func
from sqlmodel import Session, select, col
from sse_starlette import EventSourceResponse

import common
from api import R, ChatMessage, Page
from api.entity import Conversation, Message, ConversationQuery, MessageQuery
from api.enums import EventType
from api.service.conversation_service import ConversationService
from api.service.message_service import MessageService
from common import REDIS_KEY_STU_CURRENT_CONVERSATION, REDIS_KEY_CONVERSATION_AGENTS
from common.llm import LLM
from common.mysql_client import MysqlClient
from common.redis_client import RedisClient
from hk_mcp import mcp_client

router = APIRouter(prefix='/chat', tags=["chat"])
SessionDep = Annotated[Session, Depends(MysqlClient().get_session)]


@router.post('/send')
async def send(chat_msg: ChatMessage, session: SessionDep):

    if not chat_msg.stu_no:
        return R.fail(msg='学号不能为空')
    if not chat_msg.query:
        return R.fail(msg='用户输入不能为空')
    conversation_id = chat_msg.conversation_id
    llm = LLM(temperature=0.3)
    redis = RedisClient().get_aredis()
    if not conversation_id:
        if chat_msg.agents == None:
            return R.fail('必须传宝序号')
        # 新会话
        conversation = Conversation(title=llm.gen_title(chat_msg.query), stu_no=chat_msg.stu_no, agents=json.dumps(chat_msg.agents))
        ConversationService.create_conversation(conversation, session, autocommit=False)
        conversation_id = conversation.id
        await redis.set(REDIS_KEY_CONVERSATION_AGENTS.format(conversation_id), json.dumps(chat_msg.agents))

        # 插入会话记录
        prompt = PromptTemplate.from_file(f'{common.root_path()}/prompt/chat_system_prompt.md')
        feature = ''
        for agent in chat_msg.agents:
            feature += PromptTemplate.from_file(f'{common.root_path()}/prompt/feature_{agent}_prompt.md').template
        if not feature:
            feature = PromptTemplate.from_file(f'{common.root_path()}/prompt/feature_0_prompt.md').template
        sys_message = Message(conversation_id=conversation_id, role='system', content=prompt.format(stu_no=chat_msg.stu_no, current_time=datetime.now().strftime('%Y-%m-%d %H:%M:%S'), agents=chat_msg.agents, feature=feature))
        MessageService.create_messages([sys_message], session, autocommit=False)

    user_message = Message(conversation_id=conversation_id, role='human', content=chat_msg.query, tag='history')
    MessageService.create_messages([user_message], session, autocommit=False)
    session.commit()

    # 拼接历史会话
    messages = MessageService.get_messages(conversation_id, session)
    history = []
    for message in messages:
        if message.role == 'tool':
            tool = json.loads(message.content)
            history.append(('ai', tool['content']))
            # history.append(ToolMessage(content=tool['content'], tool_call_id=tool['tool_call_id'], name=tool['name']))
        else:
            history.append((message.role, message.content))

    # 设置一些标志缓存
    await redis.set(REDIS_KEY_STU_CURRENT_CONVERSATION.format(chat_msg.stu_no), conversation_id)
    # 调用MCP
    headers = {'Content-Type': 'text/event-stream'}
    # print(history)
    return EventSourceResponse(mcp_client.mcp_client(conversation_id, history, llm.qwen), headers=headers)


@router.post('/conversations')
async def conversations(cq: ConversationQuery, session: SessionDep,  current: int = 1,
                  size: Annotated[int, Query(le=100)] = 100):
    if not cq.stu_no:
        return R.fail(msg='学号不能为空')
    # 查询列表
    offset, limit = (current - 1) * size, size
    statement = MysqlClient.fill_statement(select(Conversation).offset(offset).limit(limit), Conversation, cq)
    # 如果有自定义查询字段，请在此自行添加
    statement = statement.where(Conversation.stu_no == cq.stu_no)
    kbs = session.exec(statement).all()

    # 查询总数
    count_statement = MysqlClient.fill_statement(select(func.count(col(Conversation.id))), Conversation, cq)
    count_statement = count_statement.where(Conversation.stu_no == cq.stu_no)
    total = session.exec(count_statement).one()

    page = Page.of(current=current, size=size, total=total, records=kbs)
    return R.ok(data=page)

@router.post('/messages')
async def messages(mq: MessageQuery, session: SessionDep,  current: int = 1,
                  size: Annotated[int, Query(le=100)] = 100):
    # if not mq.tu_no:
    #     return R.fail(msg='学号不能为空')
    if not mq.conversation_id:
        return R.fail(msg='会话ID不能为空')
    # 查询列表
    offset, limit = (current - 1) * size, size
    statement = MysqlClient.fill_statement(select(Message).offset(offset).limit(limit), Message, mq)
    # 如果有自定义查询字段，请在此自行添加
    statement = statement.where(Message.conversation_id == mq.conversation_id).where(Message.tag == mq.tag)
    kbs = session.exec(statement).all()

    # 查询总数
    count_statement = MysqlClient.fill_statement(select(func.count(col(Message.id))), Message, mq)
    count_statement = count_statement.where(Message.conversation_id == mq.conversation_id).where(Message.tag == mq.tag)
    total = session.exec(count_statement).one()

    page = Page.of(current=current, size=size, total=total, records=kbs)
    return R.ok(data=page)